#include <stdio.h>
#include <stdlib.h>
#include <jpeglib.h>

void nv12_to_rgb(unsigned char *yuv_data, unsigned char *rgb_data, int width, int height) {
    int frame_size = width * height;
    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            int y_index = y * width + x;
            int uv_index = frame_size + (y / 2) * width + (x & ~1); // U and V are interleaved

            int Y = yuv_data[y_index];
            int U = yuv_data[uv_index] - 128; // U is at uv_index
            int V = yuv_data[uv_index + 1] - 128; // V is at uv_index + 1

            int R = Y + 1.402 * V;
            int G = Y - 0.344136 * U - 0.714136 * V;
            int B = Y + 1.772 * U;

            // Clamp values to [0, 255]
            rgb_data[(y * width + x) * 3 + 0] = (R < 0) ? 0 : (R > 255) ? 255 : R;
            rgb_data[(y * width + x) * 3 + 1] = (G < 0) ? 0 : (G > 255) ? 255 : G;
            rgb_data[(y * width + x) * 3 + 2] = (B < 0) ? 0 : (B > 255) ? 255 : B;
        }
    }
}

void sharpen_image(unsigned char *rgb_data, int width, int height, float strength) {
    // 确保强度在合理范围内 (0.0 到 5.0)
    strength = (strength < 0.0) ? 0.0 : (strength > 5.0) ? 5.0 : strength;
    
    unsigned char *sharpened_data = (unsigned char *)malloc(width * height * 3);
    
    // 根据强度调整卷积核的值
    float center = 1.0 + 4.0 * strength;  // 中心值随强度增加
    float around = -1.0 * strength;       // 周围值随强度减小
    
    float kernel[3][3] = {
        {0, around, 0},
        {around, center, around},
        {0, around, 0}
    };

    for (int y = 1; y < height - 1; y++) {
        for (int x = 1; x < width - 1; x++) {
            float r = 0, g = 0, b = 0;

            for (int ky = -1; ky <= 1; ky++) {
                for (int kx = -1; kx <= 1; kx++) {
                    int pixel_index = ((y + ky) * width + (x + kx)) * 3;
                    r += rgb_data[pixel_index + 0] * kernel[ky + 1][kx + 1];
                    g += rgb_data[pixel_index + 1] * kernel[ky + 1][kx + 1];
                    b += rgb_data[pixel_index + 2] * kernel[ky + 1][kx + 1];
                }
            }

            // Clamp values to [0, 255]
            sharpened_data[(y * width + x) * 3 + 0] = (r < 0) ? 0 : (r > 255) ? 255 : r;
            sharpened_data[(y * width + x) * 3 + 1] = (g < 0) ? 0 : (g > 255) ? 255 : g;
            sharpened_data[(y * width + x) * 3 + 2] = (b < 0) ? 0 : (b > 255) ? 255 : b;
        }
    }

    // Copy sharpened data back to original
    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            int pixel_index = (y * width + x) * 3;
            rgb_data[pixel_index + 0] = sharpened_data[pixel_index + 0];
            rgb_data[pixel_index + 1] = sharpened_data[pixel_index + 1];
            rgb_data[pixel_index + 2] = sharpened_data[pixel_index + 2];
        }
    }

    free(sharpened_data);
}

void encode_jpeg(const char *filename, unsigned char *rgb_data, int width, int height) {
    struct jpeg_compress_struct cinfo;
    struct jpeg_error_mgr jerr;

    cinfo.err = jpeg_std_error(&jerr);
    jpeg_create_compress(&cinfo);

    FILE *outfile = fopen(filename, "wb");
    if (!outfile) {
        fprintf(stderr, "Cannot open %s\n", filename);
        return;
    }
    jpeg_stdio_dest(&cinfo, outfile);

    cinfo.image_width = width;
    cinfo.image_height = height;
    cinfo.input_components = 3; // RGB
    cinfo.in_color_space = JCS_RGB;

    jpeg_set_defaults(&cinfo);
    jpeg_start_compress(&cinfo, TRUE);

    unsigned char *row_pointer[1];
    while (cinfo.next_scanline < cinfo.image_height) {
        row_pointer[0] = &rgb_data[cinfo.next_scanline * width * 3];
        jpeg_write_scanlines(&cinfo, row_pointer, 1);
    }

    jpeg_finish_compress(&cinfo);
    fclose(outfile);
}

int main(int argc, char *argv[]) {
    if (argc < 4) {
        fprintf(stderr, "Usage: %s <input.yuv> <output.jpg> <sharpen_strength>\n", argv[0]);
        return 1;
    }

    const char *yuv_filename = argv[1];
    const char *jpeg_filename = argv[2];
    float sharpen_strength = atof(argv[3]); // Get sharpen strength from command line
    int width = 1920; // Set your image width
    int height = 1080; // Set your image height

    // Allocate memory for YUV and RGB data
    unsigned char *yuv_data = (unsigned char *)malloc(width * height * 3 / 2); // NV12 format
    unsigned char *rgb_data = (unsigned char *)malloc(width * height * 3);

    FILE *yuv_file = fopen(yuv_filename, "rb");
    if (!yuv_file) {
        fprintf(stderr, "Cannot open %s\n", yuv_filename);
        free(yuv_data);
        free(rgb_data);
        return 1;
    }
    fread(yuv_data, 1, width * height * 3 / 2, yuv_file);
    fclose(yuv_file);

    // Convert NV12 to RGB
    nv12_to_rgb(yuv_data, rgb_data, width, height);

    // Apply sharpening
    sharpen_image(rgb_data, width, height, sharpen_strength);

    // Encode to JPEG
    encode_jpeg(jpeg_filename, rgb_data, width, height);

    free(yuv_data);
    free(rgb_data);
    return 0;
} 


/*
gcc -o yuv_to_jpeg_dng yuv2jpg.c -ljpeg
*/